package com.yujiahao.bigdata.rdd.partitions

import org.apache.spark.rdd.RDD
import org.apache.spark.{Partitioner, SparkConf, SparkContext}

//TODO  自定义分区器
//TODO 实现步骤：1、写一个类继承Partitioner。2、重写方法


object Spark01_Partitions {
  def main(args: Array[String]): Unit = {
    //TODO 1、获得Spark环境的连接
    // 参数解读：setAppName("当前执行程序的名字").setMaster("当前的环境")
    // 分布式环境的基本架构为 主 - 从
    val conf: SparkConf = new SparkConf().setAppName("WordCount").setMaster("local")
    val sc: SparkContext = new SparkContext(conf)
    //TODO 3、业务逻辑
    val rdd: RDD[(String, Int)] = sc.makeRDD(List(
      ("nba", 1),
      ("cba", 2),
      ("nba", 3),
      ("wnba", 4)), 2)
    rdd.saveAsTextFile("output")
    rdd.saveAsObjectFile("output1")
    rdd.saveAsSequenceFile("output2")
    val newRDD: RDD[(String, Int)] = rdd.partitionBy(new Mypartitioner())
    newRDD.saveAsTextFile("output3")

    newRDD.collect().foreach(println)
    //TODO 2、关闭Spark环境的连接
    sc.stop()

  }

  //TODO  自定义分区器
  //TODO 实现步骤：1、写一个类继承Partitioner。2、重写方法
  class Mypartitioner extends Partitioner {

    //定义一个表示分区数量的变量
    override def numPartitions: Int = 3

    //根据数据的Key来判断数据所在的分区编号
    override def getPartition(key: Any): Int = {
      if (key == "nba") {
        0
      } else if (key == "cba") {
        1
      } else {
        2
      }
    }


  }
}
